##########################################################################################

library(data.table)
library(optparse)
library(parallel)
library(dplyr)
suppressWarnings(library(edgeR, quietly = T))
library(RColorBrewer)
library(ggplot2)
library(ggrepel)
library(clusterProfiler)
library(org.Hs.eg.db)
library(tidyverse)
library(ggsci)
library(ggrastr)

##########################################################################################

option_list <- list(
    make_option(c("--sample_list_file"), type = "character"),
    make_option(c("--rsem_file"), type = "character"),
    make_option(c("--gtf_file"), type = "character"),
    make_option(c("--out_path"), type = "character"),
    make_option(c("--foldchange_t"), type = "character"),
    make_option(c("--q_t"), type = "character"),
    make_option(c("--subtype"), type = "character")
)

if(1!=1){

    sample_list_file <- "~/20220915_gastric_multiple/dna_combinePublic/public_ref/combine/MutationInfo.combine.addMolecularSubType.rmMIX.tsv"
    rsem_file <- "~/20220915_gastric_multiple/dna_combinePublic/finalPlot/mrna/DiffGene/CombineCounts.FilterLowExpression-MergeMutiSample.TMM.tsv"
    gtf_file <- "~/ref/GTF/gencode.v19.ensg_genename.txt"
    foldchange_t <- 1.5
    q_t <- 0.05
    subtype <- "All"
    out_path <- "~/20220915_gastric_multiple/dna_combinePublic/finalPlot/revise/mrna/DiffGene/All"

}

parseobj <- OptionParser(option_list=option_list, usage = "usage: Rscript %prog [options]")
opt <- parse_args(parseobj)
print(opt)

sample_list_file <- opt$sample_list_file
rsem_file <- opt$rsem_file
out_path <- opt$out_path
gtf_file <- opt$gtf_file
q_t <- as.numeric(opt$q_t)
foldchange_t <- as.numeric(opt$foldchange_t)
subtype <- opt$subtype

dir.create(out_path , recursive = T)

##########################################################################################

info <- data.frame(fread(sample_list_file))
dat_tpm <- data.frame(fread(rsem_file))
dat_gtf <- data.frame(fread(gtf_file , header = F))
colnames(dat_gtf) <- c("gene_id" , "Hugo_Symbol")

##########################################################################################

count_norm <- dat_tpm
rownames(count_norm) <- count_norm$gene_id
count_norm <- count_norm[,-which(colnames(count_norm) == "gene_id")]

## 去除MIX的样本
use_sample <- info$Tumor[match(sapply(strsplit(colnames(count_norm), "_"), "[", 1), info$Tumor)]
count_norm <- count_norm[,which(sapply(strsplit(colnames(count_norm),"_"),"[", 1) %in% use_sample)]

##只提取全部或某个亚型的IGC和DGC样本
if(subtype=="All"){
  count_norm <- count_norm[,grep( "_IGC|_DGC|_IM", colnames(count_norm) , value = T)]
} else {
  use_sample <- info$Tumor[match(sapply(strsplit(colnames(count_norm), "_"), "[", 1), info$Tumor) & info$Molecular.subtype==subtype]
  count_norm <- count_norm[,grep( "_IGC|_DGC|_IM", colnames(count_norm) , value = T)]
  count_norm <- count_norm[,which(sapply(strsplit(colnames(count_norm),"_"),"[", 1) %in% use_sample)]
}

##########################################################################################
## 多组，两两做差异表达
conditions <- sapply(strsplit(colnames(count_norm) , "_") , "[" , 2)
levelsn <- unique(conditions)
result_diff <- c()
for( i in 1:(length(levelsn)-1) ){
    class1 <- levelsn[i]
    print(class1)
    for( j in (i+1):length(levelsn) ){
        class2 <- levelsn[j]
        print(class2)
        ## Run the Wilcoxon rank-sum test for each gene
        count_norm_use <- count_norm[,grep( paste0("_" , class1 , "|_" , class2), colnames(count_norm) , value = T)]
        conditions <- factor(sapply(strsplit(colnames(count_norm_use) , "_") , "[" , 2))

        pvalues <- sapply(1:nrow(count_norm_use),function(i){
            data<-cbind.data.frame(gene=as.numeric(t(count_norm_use[i,])),conditions)
            p <- wilcox.test(gene~conditions, data)$p.value
            return(p)
        })
        fdr=p.adjust(pvalues,method = "fdr")

        ## Calculate the fold-change for each gene
        conditionsLevel <- c(class1 , class2)
        conditions <- sapply(strsplit(colnames(count_norm_use) , "_") , "[" , 2)
        dataCon1 <- data.frame(count_norm_use[,c(which(conditions==conditionsLevel[1]))])
        dataCon2 <- data.frame(count_norm_use[,c(which(conditions==conditionsLevel[2]))])
        if(ncol(dataCon1) > 1) {
            row_means1 <- rowMeans(dataCon1)
        } else if(ncol(dataCon1) == 1) {
            row_means1 <- dataCon1[[1]]
        }
        if(ncol(dataCon2) > 1) {
            row_means2 <- rowMeans(dataCon2)
        } else if(ncol(dataCon2) == 1) {
            row_means2 <- dataCon2[[1]]
        }
        foldChanges <- log2(row_means2/row_means1)
        ## Output results based on FDR threshold
        outRst<-data.frame(log2foldChange=foldChanges, pValues=pvalues, FDR=fdr)
        rownames(outRst)=rownames(count_norm_use)

        ## class1代表分子
        outRst$class1 <- class2
        outRst$class2 <- class1
        outRst$meanTMM_class1 <- row_means2
        outRst$meanTMM_class2 <- row_means1
        outRst$gene_id <- rownames(outRst)
        outRst=na.omit(outRst)

        result_diff <- rbind( result_diff , outRst)
    }
}

colnames(result_diff) <- c("log2FoldChange" , "pvalue" , "padj" , "class1" , "class2" , "meanTMM_class1" , "meanTMM_class2" , "gene_id")

image_name <- paste0( out_path , "/DiffGene.",subtype,".tsv" )
write.table( result_diff , image_name , row.names = F , sep = "\t" , quote = F )

dat_diff <- result_diff
dat_diff <- merge( dat_diff , dat_gtf , by = "gene_id" )

##########################################################################################
## 差异表达火山图
diffplot <- function(use_dat = use_dat , image_name = image_name,subtype=subtype,Class1,Class2){

    #对原数据进行处理
    use_dat$log10_q_Value <- -log10(use_dat$padj)
    use_dat$log_FC <- use_dat$log2FoldChange
    use_dat$gene <- use_dat$Hugo_Symbol

    data <- use_dat[,c('gene','log_FC','log10_q_Value')]
    colnames(data) <- c('gene','log_FC','-log10_P_Value')

    #设置阈值
    logFC_cutoff <- log2(foldchange_t)
    log10_P_Value_cutoff <- -log10(q_t)

    options(ggrepel.max.overlaps = 20)

    plot1 <- ggplot(data = data,aes(x = log_FC,y = `-log10_P_Value`))+
      geom_point_rast(data = subset(data,abs(log_FC)<logFC_cutoff),
                 aes(size = abs(log_FC)),col = 'gray',alpha = 0.4)+
      geom_point_rast(data = subset(data,abs(`-log10_P_Value`)<log10_P_Value_cutoff & abs(log_FC)>logFC_cutoff),
                 aes(size = abs(log_FC)),col = 'gray',alpha = 0.4)+
      geom_point_rast(data = subset(data,abs(`-log10_P_Value`)>log10_P_Value_cutoff & log_FC>logFC_cutoff),
                 aes(size = abs(log_FC)),col = 'red',alpha = 0.4)+
      geom_point_rast(data = subset(data,abs(`-log10_P_Value`)>log10_P_Value_cutoff & log_FC< -logFC_cutoff),
                 aes(size = abs(log_FC)),col = 'darkgreen',alpha = 0.4)+
      theme_bw()+
      theme(legend.title = element_blank(),
            panel.grid.major = element_blank(),
            panel.grid.minor = element_blank(),
            legend.position = 'none',
            axis.line = element_line(colour = "black"))+
      labs(x='log2(fold change)',y='-log10(adjusted p-value)')+
      geom_vline(xintercept = c(-logFC_cutoff,logFC_cutoff),lty = 3,col = 'black',lwd = 0.4)+
      geom_hline(yintercept = log10_P_Value_cutoff,lty = 3,col = 'black',lwd = 0.4) +
      geom_text_repel(data = subset(data,abs(`-log10_P_Value`)>log10_P_Value_cutoff & abs(log_FC)>logFC_cutoff),
                            aes(label = gene),size = 3,col = 'black')+
      ggtitle(paste0(subtype," : ",Class1," vs ",Class2))+
       theme(plot.title = element_text(size = 16))
    ggsave( image_name , plot1 , width = 10 )
}

for( i in 1:(length(levelsn)-1) ){
    class1n <- levelsn[i]
    for( j in (i+1):length(levelsn) ){
        class2n <- levelsn[j]
        use_dat <- subset( dat_diff , class2==class1n & class1==class2n)
        image_name <- paste0( out_path , "/" , class2n , "_" , class1n , ".",subtype,".DiffGene",".pdf" )
        diffplot(use_dat = use_dat , image_name = image_name,subtype=subtype,Class2=class2n , Class1=class1n)
    }
}

##########################################################################################
## 通路富集
computGO <- function(gene_symbol = gene_symbol , type = type){
    gene_df <- bitr(gene_symbol,fromType = 'SYMBOL',toType = 'ENTREZID', 
                    OrgDb = 'org.Hs.eg.db')
    ego <- data.frame(enrichGO(gene_df$ENTREZID,OrgDb = org.Hs.eg.db,ont = 'ALL'))
    ego$type <- type
    return(ego)
}

enrich_df <- c()
for( i in 1:(length(levelsn)-1) ){
    class1n <- levelsn[i]
    for( j in (i+1):length(levelsn) ){
        class2n <- levelsn[j]

        result_diff <- subset( dat_diff , class2==class1n & class1==class2n)

        result_diff$CLUSTER <- ifelse(result_diff$log2FoldChange > log2(foldchange_t) & result_diff$padj < q_t , paste0( class2n , "_high_expression"),
          ifelse(result_diff$log2FoldChange < -log2(foldchange_t) & result_diff$padj < q_t , paste0( class1n , "_high_expression"),"no_sig"))
        result_diff <- subset(result_diff,CLUSTER %in% c(paste0( class1n , "_high_expression"),paste0( class2n , "_high_expression")))

        if(nrow(result_diff > 0 )){
            for( clust in unique(result_diff$CLUSTER) ){
                print(clust)
                gene_symbol <- unique(subset( result_diff , CLUSTER == clust )$Hugo_Symbol)
                tmp_go <- computGO(gene_symbol = gene_symbol , type = clust)
                tmp_go$class1_vs_class2 <- paste0( class1n , "_" , class2n )
                enrich_df <- rbind( enrich_df , tmp_go )
            }
        }
    }
}

write.table(enrich_df, paste0(out_path , "/enrich.go_filt_",subtype,".txt"), sep = '\t', row.names = FALSE, quote = FALSE)

##########################################################################################
## 每种类型的通路取前5最显著的
result <- c()
top_n <- 5
for( class_type in unique(enrich_df$class1_vs_class2) ){
    for(t_type in unique(enrich_df$type)){
      for( t_path in unique(enrich_df$ONTOLOGY)){
        tmp <- subset( enrich_df , type == t_type & ONTOLOGY == t_path & class1_vs_class2 == class_type )
        tmp <- tmp[order(tmp$qvalue , decreasing=F),][1:top_n,]
        result <- rbind(result , tmp)
      }
    }
}

## 画图
plotBar <- function(result = result , typeN = typeN){
  result_tmp <- subset( result , type == typeN )
  result_tmp$log10Pvalue <- -log10(result_tmp$pvalue)

  p <- ggplot(result_tmp) +
    labs(title = typeN) +
    facet_grid(.~ONTOLOGY , scales = "free_x") +
    aes(x = Description, y = log10Pvalue, fill = ONTOLOGY) +
    geom_bar(stat = "identity",colour="black") +
    scale_fill_npg()+
    ylab("-log10Pvalue") +
    theme(
      axis.title=element_text(size=15,face="plain",color="black"),
      axis.text = element_text(size=12,face="plain",color="black"),
      axis.text.x = element_text(angle = 90,colour = "black",hjust=1,vjust=0.6),
      axis.title.x = element_blank(),
      legend.postion = "none" ,
      legend.background = element_blank(),
      panel.background = element_rect(fill = "transparent",colour = "black"),
      plot.background = element_blank()
    )

  return(p)
}

for( class_type in unique(enrich_df$class1_vs_class2) ){
    result_tmp <- subset( result , class1_vs_class2 == class_type )

    if(class_type == "IM_DGC"){
        typeN <- unique(result_tmp$type)[2]
        p1 <- plotBar(result = result_tmp , typeN = typeN)
        typeN <- unique(result_tmp$type)[1]
        p2 <- plotBar(result = result_tmp , typeN = typeN)
    }else{
        typeN <- unique(result_tmp$type)[1]
        p1 <- plotBar(result = result_tmp , typeN = typeN)
        typeN <- unique(result_tmp$type)[2]
        p2 <- plotBar(result = result_tmp , typeN = typeN)
    }
    
    p <- p1 + p2

    out_file <- paste0( out_path , "/diffexp." , class_type , ".gobp." , foldchange_t , ".pdf" )
    ggsave(filename = out_file, plot = p, width = 15, height = 10)
}